import numpy as np
import cantera as ct
import multiprocessing as mp
from pathlib import Path
from multiprocessing import Pool

########################################################
chem = "Okafor2018_s59r356.yaml"
init_dataset_path = Path("combined_data.npy")
save_path = init_dataset_path.name.split('.')[0] + "_random.npy"
#init_dataset_path = Path("combined_data_10w.npy")
#save_path = init_dataset_path.name.split('.')[0] + "_random.npy"

fuel = 'NH3:0.6,CH4:0.4'
air = 'O2:1,N2:3.76'

n_process = 100
time_step = 1e-6

alpha   = 0.15

# max_C_H = 0.11764706 * 1.4
# min_C_H = 0.11764706 * 0.6

# max_N_O = 4.12 * 1.1
# min_N_O = 3.832 * 0.9

gas = ct.Solution(chem)
n_species = gas.n_species
dim_names = ['T', 'P'] + gas.species_names

idx_dict = {name: i for i, name in enumerate(dim_names)}

N2_idx = idx_dict['N2']
AR_idx = idx_dict['AR']
initial_list = list(range(2, len(dim_names)))
main_specie_list = [ 
    AR_idx,
    N2_idx,
    idx_dict['H2'],
    idx_dict['O2'],
    idx_dict['H2O'],
    idx_dict['CH4'], 
    idx_dict['CO'], 
    idx_dict['CO2'],
    idx_dict['NH3'],
    idx_dict['NO'],
    idx_dict['NO2'],
    idx_dict['N2O'],
    ]
radical_specie_list = [i for i in initial_list if i not in main_specie_list]
# print(f"Filtered List: {radical_specie_list}")
# print([dim_names[index] for index in radical_specie_list])

def BCT(x, lam = 0.1):
    return np.log(x) if lam == 0 else (np.power(x, lam) - 1)/lam

def rev_BCT(x, lam = 0.1):
    return np.exp(x) if lam == 0 else np.power(x*lam + 1, 1/lam)

def process_row(
    row, gas,    
    dim_maxs, dim_mins, 
    alpha=alpha,
    N2_idx=N2_idx, AR_idx=AR_idx, 
    radical_specie_list=radical_specie_list,
    main_specie_list=main_specie_list,
    idx_dict=idx_dict,
    ):
    
    orig_row = np.copy(row)
    BCT_row = BCT(row)
    
    k = 0
    condition = False
    while not condition and k <= 20:
        k = k +1

        maxT = dim_maxs[0]
        minT = dim_mins[0]
        maxP = dim_maxs[1]
        minP = dim_mins[1]
        
        ########################################################
        # Random perturbation
        ########################################################
        # Temperature
        # row[0] = orig_row[0] + (2*random_number-1)*100
        if BCT_row[idx_dict['H2']] < -7.5:
            upper_T, lower_T = 40, -10
        elif orig_row[0] < 320:
            upper_T, lower_T = 100, -10
        else:
            upper_T, lower_T = 200, -200
        temp_T = orig_row[0] + (np.random.random()*(upper_T - lower_T) + lower_T)
        row[0] = np.clip(temp_T, 300, 2250)
        ########################################################
        # Pressure
        row[1] = orig_row[1] + (maxP - minP)*(2*np.random.random() - 1.0)*alpha
        ########################################################
        # Species        
        # temp_ = BCT_row[N2_idx] + (np.random.random()*0.015-0.003)
        # row[N2_idx] = rev_BCT(temp_)
        
        for i in main_specie_list[1:]:
            if orig_row[0] < 350:
                upper_bnd, lower_bnd = 0.01, -0.01
            elif (i == idx_dict['N2']):
                if (orig_row[0] > 1700):
                    upper_bnd, lower_bnd = 0.0075, -0.005
                else:
                    upper_bnd, lower_bnd = 0.005, -0.005
                    
            elif i == idx_dict['H2']:
                upper_bnd, lower_bnd = 0.5, -1.
            # elif i == idx_dict['O2']:
            #     upper_bnd, lower_bnd = 0.25, -0.025
            # elif i == idx_dict['CH4']:
            #     upper_bnd, lower_bnd = 1., -0.25
            elif i == idx_dict['CO']:
                upper_bnd, lower_bnd = 0.5, -1.0  
                if 1000 < orig_row[0] < 1800:
                    lower_bnd = -1.5  
            elif i == idx_dict['CO2']:
                upper_bnd, lower_bnd = 0.5, -0.5 
            elif i == idx_dict['NO']:
                upper_bnd, lower_bnd = 0.5, -0.5
                if 1600 < orig_row[0] < 2000:
                    lower_bnd = -1.5  
            elif i == idx_dict['NO2']:
                upper_bnd, lower_bnd = 0.5, -0.5
                if 600 < orig_row[0] < 1500:
                    lower_bnd = -1. 
            elif i == idx_dict['N2O']:
                upper_bnd, lower_bnd = 0.5, -0.5
                if 500 < orig_row[0] < 1800:
                    lower_bnd = -1.2 
            
            else:
                upper_bnd, lower_bnd = 0.025, -0.025
            
            temp_ = BCT_row[i] + (np.random.random()*(upper_bnd - lower_bnd) + lower_bnd)
            temp_ = np.clip(temp_, -10, None)
            row[i] = rev_BCT(temp_)
        
        main_specie_sum = np.sum(row[main_specie_list])
        if main_specie_sum >= 1:
            continue
        
        for i in radical_specie_list:
            # temp_ = BCT_row[i] + 0.25*(BCT(dim_maxs[i]) - BCT(dim_mins[i]))*(2*np.random.random() - 1.0)
            # temp_ = np.clip(temp_, -10, None)
            # row[i] = rev_BCT(temp_)
            upper_bnd, lower_bnd = 0.9, 1.1
            if orig_row[0] < 1800:
                lower_bnd = 2.5
            
            row[i] = np.abs(orig_row[i])**(np.random.random()*(upper_bnd - lower_bnd) + lower_bnd)
        
        # row[2: -1] = row[2:-1]/np.sum(row[2:-1])*(1 - row[-1])
        row[radical_specie_list] = row[radical_specie_list]/np.sum(row[radical_specie_list])*(1 - main_specie_sum)
        ########################################################
        
        gas.TPY = row[0], row[1], row[2:] 

        # H_mole_fraction = gas.elemental_mole_fraction("H")
        # O_mole_fraction = gas.elemental_mole_fraction("O")
        # N_mole_fraction = gas.elemental_mole_fraction("N")
        # C_mole_fraction = gas.elemental_mole_fraction("C")  
        
        # C_H_ratio = C_mole_fraction / H_mole_fraction
        # N_O_ratio = N_mole_fraction / O_mole_fraction
        
        equiv_ratio = gas.equivalence_ratio(fuel, air)
        
        condition = (
            #     min_C_H <= C_H_ratio <= max_C_H
            # and min_N_O <= N_O_ratio <= max_N_O
                dim_mins[N2_idx]-(dim_maxs[N2_idx]-dim_mins[N2_idx])*0.04 <= row[N2_idx] # <= dim_maxs[N2_idx]+(dim_maxs[N2_idx]-dim_mins[N2_idx])*0.06
            and 300 <= row[0] <= (maxT + 200)
            and 0.95 <= equiv_ratio <= 1.05
            and main_specie_sum <= 1
        )
    
    return row if condition else None

if __name__ == '__main__':
    
    print("Loading Original Sample Data...")
    test = np.load(init_dataset_path)
    print(f"Shape of Original Sample Data:    {test.shape}")
    print("="*50)

    ########################################################
    # Data augmentation
    ########################################################
    # test[:, 2:] = BCT(test[:, 2:])
    
    dim_maxs = np.max(test, axis=0)
    dim_mins = np.min(test, axis=0)
    
    for i, dim_name in enumerate(dim_names):
        print(f"{dim_name:<8} {dim_maxs[i]:<12.5e} - {dim_mins[i]:<12.5e}")

    r_n = 500 # random number
    test_tmp = np.copy(test)
    rows_per_process = (test_tmp.shape[0]+n_process-1)//n_process

    print(f"\nFiltered Array Shape: {test_tmp.shape=}\n")
    test_r = np.empty((0, test_tmp.shape[1]), float)

    def process_batch(process_num):
        results_process = np.empty((0, test_tmp.shape[1]), float)
        for n in range(process_num*rows_per_process, min((process_num+1)*rows_per_process, test_tmp.shape[0])):
            random_row = process_row(test_tmp[n], gas, dim_maxs, dim_mins)
            if random_row is not None:
                results_process = np.vstack((results_process, random_row))
        
        return results_process

    for count in range(r_n):
        print(count)

        with Pool(processes=n_process) as pool:
            results = pool.map(process_batch, range(n_process))
            # print([result.shape for result in results])
            
            test_r = np.vstack((test_r, np.vstack(results)))
            print(test_r.shape)
            
                
        if test_r.shape[0] > 7500000: break

    np.random.shuffle(test_r)

    print(f"\nFiltered Array Shape: {test_r.shape=}\n")
    res_dim_maxs = np.max(test_r, axis=0)
    res_dim_mins = np.min(test_r, axis=0)
    
    for i, dim_name in enumerate(dim_names):
        print(f"{dim_name:<8} {res_dim_maxs[i]:<12.5e} - {res_dim_mins[i]:<12.5e}")

    np.save(save_path, test_r)
    print(f"Saved as {save_path}")

    exit()
